import asyncio
import websockets
import numpy as np
import threading
import socket
import cv2
import traceback

# 全局变量存储编码后的JPEG数据和客户端集合
latest_jpeg = None
jpeg_lock = threading.Lock()
connected_clients = set()

# 摄像头捕获线程函数
def camera_capture():
    global latest_jpeg
    cap = cv2.VideoCapture(0)  # 0表示默认摄像头
    if not cap.isOpened():
        raise RuntimeError("无法打开摄像头")
    
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            # 单次编码（关键优化点）
            ret, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
            if ret:
                with jpeg_lock:
                    latest_jpeg = jpeg.tobytes()
                    
            cv2.waitKey(1)  # 避免阻塞
    finally:
        cap.release()

# 广播协程（集中发送数据）
async def broadcaster():
    while True:
        with jpeg_lock:
            if latest_jpeg:
                # 批量发送给所有客户端
                await asyncio.gather(*[
                    client.send(latest_jpeg) 
                    for client in connected_clients
                ], return_exceptions=True)
        await asyncio.sleep(1/20)  # 控制20 FPS

# 客户端连接处理
async def client_handler(websocket):
    connected_clients.add(websocket)
    try:
        await websocket.wait_closed()
    finally:
        connected_clients.remove(websocket)

# 获取本地IP地址
def get_local_ip():
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("8.8.8.8", 80))
            return s.getsockname()[0]
    except:
        return "127.0.0.1"

async def main():
    # 启动摄像头捕获线程
    camera_thread = threading.Thread(target=camera_capture, daemon=True)
    camera_thread.start()

    # 等待摄像头初始化
    while True:
        with jpeg_lock:
            if latest_jpeg is not None:
                break
        await asyncio.sleep(0.1)

    # 启动广播协程
    asyncio.create_task(broadcaster())

    # 启动WebSocket服务器
    local_ip = get_local_ip()
    print(f"WebSocket服务器运行在 {local_ip}:3001")
    
    async with websockets.serve(client_handler, "0.0.0.0", 3001):
        print("WebSocket服务器已启动")
        await asyncio.Future()  # 持续运行

if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("程序终止")
    finally:
        cv2.destroyAllWindows()